import os
import joblib
import numpy as np
from PIL import Image
import gradio as gr


def to_gray_flat_pil(img: Image.Image, size: int) -> np.ndarray:
	arr = img.convert("L").resize((size, size))
	x = np.asarray(arr, dtype=np.float32) / 255.0
	return x.reshape(1, -1)


def load_payload_any(base_dir: str):
	candidates = [
		os.path.join(base_dir, "best_ensemble.pkl"),
		os.path.join(base_dir, "models", "best_ensemble.pkl"),
	]
	for p in candidates:
		if os.path.exists(p):
			return joblib.load(p), p
	raise FileNotFoundError("未找到模型文件：尝试了 best_ensemble.pkl 以及 models/best_ensemble.pkl")


def build_app():
	base = os.path.dirname(os.path.abspath(__file__))
	payload, model_path = load_payload_any(base)
	model = payload["model"]
	class_names = payload.get("class_names", ["cat", "dog"])  # index: 0->cat, 1->dog
	img_size = int(payload.get("img_size", 32))

	def infer(img: Image.Image):
		x = to_gray_flat_pil(img, img_size)
		if hasattr(model, "predict_proba"):
			proba = model.predict_proba(x)[0]
		else:
			# 若无概率输出，用预测标签构造 one-hot 概率
			pred = int(model.predict(x)[0])
			proba = np.zeros(len(class_names), dtype=np.float32)
			proba[pred] = 1.0

		idx = int(np.argmax(proba))
		scores = {class_names[i]: float(proba[i]) for i in range(len(class_names))}
		label_cn = "猫" if class_names[idx] == "cat" else "狗"
		# 返回两路输出：文本标签 + 概率分布（gr.Label 支持 dict 显示）
		return label_cn, scores

	with gr.Blocks(title="猫狗分类（集成学习）") as demo:
		gr.Markdown("# 猫狗分类（集成学习）\n上传一张图片，模型将预测猫或狗，并给出概率。")
		with gr.Row():
			inp = gr.Image(type="pil", label="上传图片")
			out_label = gr.Textbox(label="预测")
			out_scores = gr.Label(label="概率分布")
		btn = gr.Button("预测")
		btn.click(infer, inputs=inp, outputs=[out_label, out_scores])
		gr.Markdown(f"当前加载模型路径：`{os.path.relpath(model_path, base)}`")

	return demo


if __name__ == "__main__":
	app = build_app()
	# 默认本地启动；如本机无法访问 localhost，可设置 share=True
	app.launch(share=True)

